extern crate alloc;

use alloc::alloc::{alloc as std_alloc, dealloc as std_dealloc, Layout};

// Count of the current number of allocated bytes.
#[allow(non_upper_case_globals)]
static mut allocated_bytes: usize = 0;

// Limit on number of allocations that are possible.
#[allow(non_upper_case_globals)]
static mut allocation_limit: Option<usize> = None;

/// Set an artificial limit on the amount of memory that can be allocated.
pub unsafe fn set_limit(alloc_count: usize) {
    unsafe {
        allocation_limit = Some(alloc_count);
    }
}

/// Get a count of the number of bytes currently allocated.
pub unsafe fn get_allocated() -> usize {
    unsafe { allocated_bytes }
}

// Overwrite a block of memory with a repeated pattern.
/* unsafe fn overwrite(bytes: *mut [u8], pattern: usize) {
    // split the pattern into bytes
    let pattern_bytes = [
        (pattern >> 24) as u8,
        (pattern >> 16) as u8,
        (pattern >> 8) as u8,
        pattern as u8,
    ];

    // Fill the memory with the pattern
    for i in 0..bytes.len() {
        let pattern_byte = pattern_bytes[i % 4];
        unsafe {
            let ptr = bytes.cast::<u8>().add(i);
            core::ptr::write(ptr, pattern_byte);
        }
    }
}
 */
/// Allocate a block of memory.
pub unsafe fn malloc(layout: Layout) -> *mut u8 {
    let bytes = layout.size();

    // Check if we have reached the allocation limit.
    if unsafe { allocation_limit } == Some(0) {
        return core::ptr::null_mut();
    }

    let ptr = unsafe { std_alloc(layout) };

    // Update counter
    unsafe {
        allocated_bytes += bytes;
    }

    // Decrease the allocation limit
    #[allow(static_mut_refs)]
    if let Some(limits) = unsafe { allocation_limit.as_mut() } {
        if *limits > 0 {
            *limits -= 1;
        } else {
            unsafe {
                allocation_limit = None;
            }
        }
    }

    ptr
}

/// Deallocate a block of memory.
pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
    // If the pointer is null, do nothing.
    if ptr.is_null() {
        return;
    }

    unsafe {
        std_dealloc(ptr, layout);
    }

    // Update counter
    unsafe {
        allocated_bytes -= layout.size();
    }
}

/// Reallocate a previously-allocated block to a new size, preserving contents.
pub unsafe fn realloc(ptr: *mut u8, layout: Layout, bytes: usize) -> *mut u8 {
    let new_layout = Layout::from_size_align(bytes, layout.align()).unwrap();
    let new_ptr = unsafe { malloc(new_layout) };

    if new_ptr.is_null() {
        return new_ptr;
    }

    if !ptr.is_null() {
        let bytes_to_copy = core::cmp::min(layout.size(), new_layout.size());

        unsafe { core::ptr::copy_nonoverlapping(ptr.cast_const(), new_ptr, bytes_to_copy) };

        unsafe { dealloc(ptr, layout) };
    }

    new_ptr
}

/// Allocate a block of memory for an array of structures, initialising
/// the contents to zero.

pub unsafe fn calloc(layout: Layout) -> *mut u8 {
    let result = unsafe { malloc(layout) };

    if result.is_null() {
        return result;
    }

    core::ptr::write_bytes(result, 0, layout.size());
    result
}
